import os
from ultralytics import YOLO
from matplotlib import rcParams
from matplotlib.font_manager import FontProperties

# Handle potential library conflicts
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# Set up Chinese font support for matplotlib
#font_path = './simhei.ttf'
#if not os.path.exists(font_path):
#    print("Error: SimHei font file not found at the specified path.")
#else:
#    prop = FontProperties(fname=font_path)
#    rcParams['font.family'] = prop.get_name()

def main():
    # Model path and dataset file
    model_path = 'yolo11n.pt'
    data_path = './ultralytics/cfg/datasets/coco.yaml'
    #data_path = './datasets/data.yaml'
    # Load YOLO model
    model = YOLO(model_path)

    # Train the model with enhanced preprocessing for lighting and partial target recognition
    results = model.train(
        data=data_path,
        epochs=50,
        imgsz=640,  # Adjust this value to experiment with different image sizes
        batch=32,
        name='exp',
        workers=4,
        verbose=True,
        augment=True,          # Enable data augmentation
        mosaic=0.4,            # Increased mosaic for better handling of different object sizes
        hsv_h=0.5,             # Moderate hue shift for color diversity
        fliplr=0.5,            # Horizontal flip to enhance diversity
        flipud=0.2,            # Minor vertical flip to simulate slight orientation variance
        scale=0.5,             # Introduced scaling for target size variation
        patience=0,          # Early stopping patience to allow ample training
        lr0=0.01,              # Initial learning rate
        lrf=0.00001            # Final learning rate for gradual decay
    )

if __name__ == '__main__':
    main()

